import ptan
import numpy as np
from typing import List, Optional

import torch
import torch.nn as nn
import torch.nn.functional as F

from lib import common

HID_SIZE = 128


'''
在强化学习的Actor-Critic方法中，特别是在处理连续动作空间时，网络通常预测动作的均值（`mu`）和方差（`var`）。在你给出的代码中，模型使用不同的激活函数来处理这两个输出，即`tanh`用于均值，而`Softplus`用于方差。让我们分别来看看这两个激活函数的选择原因：

### 1. 均值 (`mu`) 的 `tanh` 激活函数

- **动作范围限制**：`tanh` 函数的输出范围在 [-1, 1] 之间。在许多连续动作空间的问题中，动作的有效范围被限制在这个区间内。例如，如果你在一个物理模拟环境中控制机械臂，动作可能是旋转角度，这些角度通常有最大和最小值。

- **输出规范化**：使用 `tanh` 确保网络输出的均值永远不会超出这个范围，这样可以保证生成的动作始终在有效的动作空间内。

### 2. 方差 (`var`) 的 `Softplus` 激活函数

- **保证正值**：方差是衡量数据离散程度的度量，它必须是正值。`Softplus` 函数是一个平滑的函数，其输出始终为正值。`Softplus` 函数的公式是 `log(1 + exp(x))`，对于所有的输入 `x`，输出都是正的。

- **灵活性**：与 `ReLU` 函数相比，`Softplus` 在 `x` 接近0时更加平滑，这允

许更精细的控制方差的值。`Softplus` 提供了一种有效的方式来表示方差，即使在其原始预测值接近零或负值时也能保持正数。这在学习过程中非常有用，特别是当网络还未充分训练，对方差的估计可能接近零或变化较大时。

### 结合使用 `tanh` 和 `Softplus`

在这个A2C模型中，`tanh` 用于输出规范化到特定范围内的动作均值，而 `Softplus` 用于确保方差始终为正值且平滑。这种设计使模型能够有效地操作连续动作空间，同时保持必要的灵活性和稳定性。使用这两个激活函数的组合，模型可以生成符合实际物理限制的动作（如限制在一定范围内的力或速度），并且以合理的方式探索这些动作的不确定性。
'''
class ModelA2C(nn.Module):
    def __init__(self, obs_size, act_size):
        '''
        obs_size: 观测的数据维度，不是图像数据，所以后续用的是全连接层
        act_size: 动作空间的维度，在这个游戏里面，指的是同时执行动作的数量
        '''

        super(ModelA2C, self).__init__()

        # 只有这个是提取特征，其余的都是输出结果
        self.base = nn.Sequential(
            nn.Linear(obs_size, HID_SIZE),
            nn.ReLU(),
        )
        # 输出均值，表示最终要执行的动作内容
        # 这里由于预测的输出动作包含负值，所以采用tanh函数，将输出值限制在-1到1之间
        # 而不是用sigmoid
        self.mu = nn.Sequential(
            nn.Linear(HID_SIZE, act_size),
            nn.Tanh(),
        )
        # https: // zhuanlan.zhihu.com / p / 461707201
        # var作用 方差平方，怀疑是用来稳定输出的概率范围大小
        # 用来促进网络进行探索以及指导网路朝哪个方向进行训练，使得整体趋近中值，但是在这里并没有直接使用方差，而是使用了信息熵的方式
        self.var = nn.Sequential(
            nn.Linear(HID_SIZE, act_size),
            nn.Softplus(), # Relu的替代函数，用于解决梯度消失问题 具体使用场景查看笔记内容
        )
        # 状态值（Q值），用来评价当前Q值，来评估当前执行的动作是否有优势
        self.value = nn.Linear(HID_SIZE, 1)

    def forward(self, x):
        '''
        return 均值，方差平方，Q值
        '''
        base_out = self.base(x)
        return self.mu(base_out), self.var(base_out), self.value(base_out)


class DDPGActor(nn.Module):
    '''
    深度确定性策略梯度动作预测网络
    '''
    def __init__(self, obs_size, act_size):
        '''
        obs_size: 环境的维度
        act_size: 能够同时执行动作的个数（比如有多个手，不是每个手可以执行哪些动作）
        '''
        super(DDPGActor, self).__init__()

        obs_action = (obs_size[2], obs_size[0], obs_size[1])
        self.conv = nn.Sequential(
            nn.Conv2d(obs_action[0], 32, kernel_size=8, stride=4),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1),
            nn.BatchNorm2d(64),
            nn.ReLU()
        )

        conv_out_size = self._get_conv_out(obs_action)
        self.fc = nn.Sequential(
            nn.Linear(conv_out_size, 512),
            nn.ReLU(),
            nn.Linear(512, act_size),
            nn.Tanh()
        )

    def _get_conv_out(self, shape):
        o = self.conv(torch.zeros(1, *shape))
        return int(np.prod(o.size()))

    def forward(self, x):
        fx = x.float() / 256
        conv_out = self.conv(fx).view(fx.size()[0], -1)
        return self.fc(conv_out)


class DDPGCritic(nn.Module):
    '''
    深度确定性策略梯度网络Q值评价网络
    '''
    def __init__(self, obs_size, act_size):
        '''
        obs_size: 环境的维度
        act_size: 能够同时执行动作的个数（比如有多个手，不是每个手可以执行哪些动作）
        '''
        super(DDPGCritic, self).__init__()

        obs_action = (obs_size[2], obs_size[0], obs_size[1])
        self.conv = nn.Sequential(
            nn.Conv2d(obs_action[0], 32, kernel_size=8, stride=4),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1),
            nn.BatchNorm2d(64),
            nn.ReLU()
        )

        conv_out_size = self._get_conv_out(obs_action)
        self.fc = nn.Sequential(
            nn.Linear(conv_out_size + act_size, 512),
            nn.ReLU(),
            nn.Linear(512, 1)
        )

    def _get_conv_out(self, shape):
        o = self.conv(torch.zeros(1, *shape))
        return int(np.prod(o.size()))

    def forward(self, x, a):
        fx = x.float() / 256
        conv_out = self.conv(fx).view(fx.size()[0], -1)
        return self.fc(torch.cat([conv_out, a], dim=1))


# class DDPGActorSimple(nn.Module):
#     '''
#     深度确定性策略梯度动作预测网络
#     '''
#     def __init__(self, obs_size, act_size):
#         '''
#         obs_size: 环境的维度
#         act_size: 能够同时执行动作的个数（比如有多个手，不是每个手可以执行哪些动作）
#         '''
#         super(DDPGActorSimple, self).__init__()
#         self.train_action = True
#         self.train_qvalue = True

#         obs_action = (obs_size[2], obs_size[0], obs_size[1])
#         self.conv = nn.Sequential(
#             nn.Conv2d(obs_action[0], 32, kernel_size=8, stride=4),
#             nn.BatchNorm2d(32),
#             nn.ReLU(),
#             nn.Conv2d(32, 64, kernel_size=4, stride=2),
#             nn.BatchNorm2d(64),
#             nn.ReLU(),
#             nn.Conv2d(64, 64, kernel_size=3, stride=1),
#             nn.BatchNorm2d(64),
#             nn.ReLU()
#         )

#         conv_out_size = self._get_conv_out(obs_action)
#         self.action = nn.Sequential(
#             nn.Linear(conv_out_size, 512),
#             nn.ReLU(),
#             nn.Linear(512, act_size),
#             nn.Tanh()
#         )

#         self.qvalue = nn.Sequential(
#             nn.Linear(conv_out_size + act_size, 512),
#             nn.ReLU(),
#             nn.Linear(512, 1)
#         )

#     def _get_conv_out(self, shape):
#         o = self.conv(torch.zeros(1, *shape))
#         return int(np.prod(o.size()))

#     def forward(self, x, a=None):
#         fx = x.float() / 256
#         conv_out = self.conv(fx).view(fx.size()[0], -1)

#         if a is None:
#             with torch.set_grad_enabled(self.train_action):
#                 pred_action = self.action(conv_out)
#             with torch.set_grad_enabled(self.train_qvalue):
#                 qvalue = self.qvalue(torch.cat([conv_out, pred_action], dim=1))
#         else:
#             with torch.set_grad_enabled(self.train_action):
#                 pred_action = self.action(conv_out)
#             with torch.set_grad_enabled(self.train_qvalue):
#                 qvalue = self.qvalue(torch.cat([conv_out, a], dim=1))
#         return pred_action, qvalue

#     def set_train_qvalue(self, freeze=True):
#         for param in self.action.parameters():
#             param.requires_grad = not freeze

#     def set_train_action(self, freeze=True):
#         for param in self.qvalue.parameters():
#             param.requires_grad = not freeze

#     def get_shared_parameter(self):
#         return list(self.conv.parameters())

#     def get_action_parameter(self):
#         return list(self.action.parameters()) + self.get_shared_parameter()

#     def get_qvalue_parameter(self):
#         return list(self.qvalue.parameters()) + self.get_shared_parameter()


class DDPGActorSimple(nn.Module):
    '''
    深度确定性策略梯度动作预测网络
    '''
    def __init__(self, obs_size, act_size):
        '''
        obs_size: 环境的维度
        act_size: 能够同时执行动作的个数（比如有多个手，不是每个手可以执行哪些动作）
        '''
        super(DDPGActorSimple, self).__init__()

        obs_action = (obs_size[2], obs_size[0], obs_size[1])
        self.conv = nn.Sequential(
            nn.Conv2d(obs_action[0], 32, kernel_size=8, stride=4),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1),
            nn.BatchNorm2d(64),
            nn.ReLU()
        )

        conv_out_size = self._get_conv_out(obs_action)
        self.action = nn.Sequential(
            nn.Linear(conv_out_size, 512),
            nn.ReLU(),
            nn.Linear(512, act_size),
            nn.Tanh()
        )

        self.qvalue = nn.Sequential(
            nn.Linear(conv_out_size + act_size, 512),
            nn.ReLU(),
            nn.Linear(512, 1)
        )

    def _get_conv_out(self, shape):
        o = self.conv(torch.zeros(1, *shape))
        return int(np.prod(o.size()))

    def forward(self, x, a=None):
        fx = x.float() / 256
        conv_out = self.conv(fx).view(fx.size()[0], -1)

        if a is None:
            pred_action = self.action(conv_out)
            qvalue = self.qvalue(torch.cat([conv_out, pred_action], dim=1))
        else:
            pred_action = self.action(conv_out)
            qvalue = self.qvalue(torch.cat([conv_out, a], dim=1))
        return pred_action, qvalue
    
    def qval_params(self):
        return list(self.qvalue.parameters()) + list(self.conv.parameters())
    

    def actor_params(self):
        return self.action.parameters()

class D4PGCritic(nn.Module):
    '''
    D4PG自己的Q值评价网路
    '''
    def __init__(self, obs_size, act_size, n_atoms, v_min, v_max):
        super(D4PGCritic, self).__init__()

        # 环境特征采样观测网络
        self.obs_net = nn.Sequential(
            nn.Linear(obs_size, 400),
            nn.ReLU(),
        )

        # 环境特征采样+动作特征采样组合特征采集网络
        # 不同之处是返回的是一个Q值的概率分布
        self.out_net = nn.Sequential(
            nn.Linear(400 + act_size, 300),
            nn.ReLU(),
            nn.Linear(300, n_atoms)
        )

        # 概率分布每个区间的大小
        delta = (v_max - v_min) / (n_atoms - 1)
        # todo supports作用
        self.register_buffer("supports", torch.arange(v_min, v_max + delta, delta))

    def forward(self, x, a):
        '''
        推理并返回预测Q值的概率分布
        '''

        obs = self.obs_net(x)
        return self.out_net(torch.cat([obs, a], dim=1))

    def distr_to_q(self, distr):
        '''
        将Q值分布转换为Q值
        '''
        # 计算概率分布，将概率分布乘以回报分布（-10到10），得到回报分布值
        weights = F.softmax(distr, dim=1) * self.supports
        # 计算加权回报分布得到期望Q值
        res = weights.sum(dim=1)
        return res.unsqueeze(dim=-1)


class AgentA2C(ptan.agent.BaseAgent):
    '''
    A2C连续策略值网络代理，之所以自己写，是因为需要应对连续值的动作输出
    '''
    def __init__(self, net, device="cpu"):
        '''
        net: 神经网路
        device: 指定网路的设备
        '''
        self.net = net
        self.device = device

    def __call__(self, states, agent_states):
        '''
        states: 观察的状态
        agent_state: 代理器的状态，在这里没有直接使用，在某些地方可能需要用到代理器的内部状态

        return: 选择执行的动作值，代理器的状态
        '''
        states_v = ptan.agent.float32_preprocessor(states).to(self.device)

        # 将观测传入网络得到均值和方差
        mu_v, var_v, _ = self.net(states_v)
        mu = mu_v.data.cpu().numpy()
        sigma = torch.sqrt(var_v).data.cpu().numpy()
        # 以下实在计算高斯分布吗？
        # 不是：以mu为均值，sigma为标准差的正态分布中抽取随机数，作为每个状态下采取的动作。在连续动作空间的问题中，这种方法能够在探索和利用之间取得平衡。随机性（通过sigma体现）允许探索不同的动作，而均值（mu）通常代表了当前网络认为最佳的动作。
        # 综上所述，这个是在从均值和方差中随机选择需要执行的动作值
        actions = np.random.normal(mu, sigma)
        actions = np.clip(actions, -1, 1)
        return actions, agent_states


class AgentAcktr(ptan.agent.BaseAgent):
    '''
    创建代理器
    '''
    def __init__(self, net, device="cpu"):
        self.net = net
        self.device = device

    def __call__(self, states, agent_states):
        '''
        states: 观测的环境状态
        agent_states：智能体自己的状态，在这里是没有使用的
        '''
        # 创建环境预处理器，将环境状态转换为float32类型
        states_v = ptan.agent.float32_preprocessor(states).to(self.device)

        # 通过环境状态预测执行的动作
        mu_v = self.net(states_v)
        mu = mu_v.data.cpu().numpy()
        logstd = self.net.logstd.data.cpu().numpy()
        # 该动作的作用，是对预测的动作添加随机噪音，实现动作的探索
        actions = mu + np.exp(logstd) * np.random.normal(size=logstd.shape)
        # 将执行的动作压缩到-1到1中，可能是因为输入给网络的值不能超过-1和1
        actions = np.clip(actions, -1, 1)
        return actions, agent_states



class AgentPPO(ptan.agent.BaseAgent):
    '''
    创建代理器
    '''
    def __init__(self, net, device="cpu"):
        self.net = net
        self.device = device

    def __call__(self, states, agent_states):
        '''
        states: 观测的环境状态
        agent_states：智能体自己的状态，在这里是没有使用的
        '''
        # 创建环境预处理器，将环境状态转换为float32类型
        states_v = ptan.agent.float32_preprocessor(states).to(self.device)

        # 通过环境状态预测执行的动作
        mu_v = self.net(states_v)
        mu = mu_v.data.cpu().numpy()
        logstd = self.net.logstd.data.cpu().numpy()
        # 该动作的作用，是对预测的动作添加随机噪音，实现动作的探索
        actions = mu + np.exp(logstd) * np.random.normal(size=logstd.shape)
        # 将执行的动作压缩到-1到1中，可能是因为输入给网络的值不能超过-1和1
        actions = np.clip(actions, -1, 1)
        return actions, agent_states

class AgentDDPGSimple(ptan.agent.BaseAgent):
    """
    Agent implementing Orstein-Uhlenbeck exploration process
    实现一个具备探索能力的智能体
    """
    def __init__(self, net, device="cpu", ou_enabled=True, ou_mu=0.0, ou_teta=0.15, ou_sigma=0.2, ou_epsilon=1.0):

        self.net = net
        self.device = device
        self.ou_enabled = ou_enabled
        self.ou_mu = ou_mu
        self.ou_teta = ou_teta
        self.ou_sigma = ou_sigma
        self.ou_epsilon = ou_epsilon

    def initial_state(self):
        # 统一接口，但是这里不需要使用，返回None而不是使用pass
        return None

    def __call__(self, states, agent_states):
        '''
        states：当前的环境状态
        agent_states: 内部智能体的状态,之前的代理器里面这个基本无用,一开始的时候，agent_states是空的，但在这里因为要使用OU过程对执行的动作进行噪音干扰，所以需要使用了智能体的内部状态
        '''
        # 将环境转换为目标的数据类型
        states_v = ptan.agent.float32_preprocessor(states).to(self.device)
        # 得到执行的动作输出
        mu_v, _ = self.net(states_v)
        actions = mu_v.data.cpu().numpy()

        if self.ou_enabled and self.ou_epsilon > 0:
            # 启动随机探索
            new_a_states = []
            for a_state, action in zip(agent_states, actions):
                if a_state is None:
                    # 为什么智能体状态为空时可以构建一个维度和动作相同，且全为零的状态
                    # 这是针对初始状态，初始情况下，智能体的状态为空，则创建一个为0的状态
                    a_state = np.zeros(shape=action.shape, dtype=np.float32)
                # 书p298也 todo 了解OU过程
                a_state += self.ou_teta * (self.ou_mu - a_state)
                a_state += self.ou_sigma * np.random.normal(size=action.shape)

                action += self.ou_epsilon * a_state
                new_a_states.append(a_state)
        else:
            new_a_states = agent_states

        # 这个步骤 的意思是修正动作值到-1和1之间，否则将导致PyBullet抛出异常
        actions = np.clip(actions, -1, 1)
        return actions, new_a_states


class AgentDirect(ptan.agent.BaseAgent):
    """
    Agent implementing Orstein-Uhlenbeck exploration process
    实现一个具备探索能力的智能体
    """
    def __init__(self, net, device="cpu", ):

        self.net = net
        self.device = device

    def initial_state(self):
        # 统一接口，但是这里不需要使用，返回None而不是使用pass
        return None

    def __call__(self, states, agent_states):
        '''
        states：当前的环境状态
        agent_states: 内部智能体的状态,之前的代理器里面这个基本无用,一开始的时候，agent_states是空的，但在这里因为要使用OU过程对执行的动作进行噪音干扰，所以需要使用了智能体的内部状态
        '''
        # 将环境转换为目标的数据类型
        states_v = ptan.agent.float32_preprocessor(states).to(self.device)
        # 得到执行的动作输出
        mu_v = self.net(states_v)
        actions = mu_v.data.cpu().numpy()
        new_a_states = agent_states

        # 这个步骤 的意思是修正动作值到-1和1之间，否则将导致PyBullet抛出异常
        actions = np.clip(actions, -1, 1)
        return actions, new_a_states


class AgentDDPG(ptan.agent.BaseAgent):
    """
    Agent implementing Orstein-Uhlenbeck exploration process
    实现一个具备探索能力的智能体
    """
    def __init__(self, net, device="cpu", ou_enabled=True, ou_mu=0.0, ou_teta=0.15, ou_sigma=0.2, ou_epsilon=1.0):
        '''
        net: 动作执行网络
        device: 执行的设备
        ou_enabled: 是否开启OU过程探索
        ou_mu: 查看OU过程说明
        ou_teta: 查看OU过程说明
        ou_sigma: 查看OU过程说明
        ou_epsilon: 用于控制OU过程对最终动作的执行影响大小

        Ornstein-Uhlenbeck（OU）过程是一个用来生成噪声的数学工具，特别适合用于模拟物理过程中存在的摩擦或阻力的场景，因此它在强化学习中被用作探索机制，尤其是在连续动作空间的问题中。

        在使用OU过程时，参数`ou_teta`、`ou_mu`、`ou_sigma`通常是经验性选择的，他们分别控制：

        1. **`ou_teta`（Theta）** - 这个参数决定了系统回到均值（或趋向平稳状态）的速度。较高的`ou_teta`值会使系统更快地返回到均值，即产生的噪声会更快地趋于稳定状态。如果`ou_teta`设置得太高，系统可能会过于迅速地稳定下来，从而减少探索。太低，则可能导致系统的探索过于缓慢和延迟。

        2. **`ou_mu`（Mu）** - 这是长期平均或均值项，OU过程将围绕此均值进行波动。在强化学习的上下文中，`ou_mu`通常设置为0，这意味着没有任何行动的偏好，确保探索不会偏向于任何特定的行动。

        3. **`ou_sigma`（Sigma）** - 这个参数代表了波动的振幅或标准差，决定了噪声的强度或大小。较大的`ou_sigma`增加了探索的范围和随机性，但可能导致智能体行为变得不稳定；较小的值将减少探索性噪声，可能导致智能体过早收敛到次优策略。

        在确定这些参数时，通常需要通过实验调整以找到适合特定任务的值。一些普遍的做法和建议包括：

        - 开始时使用中等范围的值进行初步测试。
        - 观察智能体的表现，如果它行动过于随机，可以降低`ou_sigma`。
        - 如果智能体似乎探索得不够，或者过早收敛到一个行动上，可以增加`ou_sigma`。
        - 调整`ou_teta`以控制噪声的“记忆”效果，即噪声持续影响行动的程度。

        调整这些参数通常需要反复试验，因为理想的值很大程度上取决于具体问题、智能体的体系结构、奖励结构以及训练过程的其他方面。此外，有时会随着训练的进展而逐步减少噪声的影响，这是通过减少`ou_epsilon`来实现的，这种做法类似于退火。

        OU过程中的参数值ou_mu、ou_theta、ou_sigma和ou_dt通常根据具体问题来确定,没有一个通用的设置方法。下面是一些确定这些参数的常见思考:

        1. ou_mu - 长期均值可以设置为目标值或期望的平均奖励值。

        2. ou_theta - 回归速度可以先设置为一个较小的值(例如0.15),表示轻度回归到均值。可以通过调参来获得更好的回归速度。

        3. ou_sigma - 随机扰动标准差可以先设置较小(例如0.3),然后通过调参确定合适的探索强度。

        4. ou_dt - 时间步长可以设置得较小(例如0.01),然后确保dt足够小以维持迭代稳定性。

        5. ou_x0 - 通常初始化为0即可。

        另外,这些参数还需要互相平衡。例如theta越大,则需要更大的sigma来增加探索。

        一个通用的思路是:

        - 先用较小的回归速度θ,较大的扰动σ开始,更偏向探索。

        - 然后观察过程的变化,如果变化太随机可以适当增大θ或减小σ。

        - 调节参数直到得到既能利用已知信息又能适度探索的平衡效果。

        所以确定OU过程的参数需要根据具体问题,通过试错来找到一个合适的组合。
        '''
        self.net = net
        self.device = device
        self.ou_enabled = ou_enabled
        self.ou_mu = ou_mu
        self.ou_teta = ou_teta
        self.ou_sigma = ou_sigma
        self.ou_epsilon = ou_epsilon

    def initial_state(self):
        # 统一接口，但是这里不需要使用，返回None而不是使用pass
        return None

    def __call__(self, states, agent_states):
        '''
        states：当前的环境状态
        agent_states: 内部智能体的状态,之前的代理器里面这个基本无用,一开始的时候，agent_states是空的，但在这里因为要使用OU过程对执行的动作进行噪音干扰，所以需要使用了智能体的内部状态
        '''
        # 将环境转换为目标的数据类型
        states_v = ptan.agent.float32_preprocessor(states).to(self.device)
        # 得到执行的动作输出
        mu_v = self.net(states_v)
        actions = mu_v.data.cpu().numpy()

        if self.ou_enabled and self.ou_epsilon > 0:
            # 启动随机探索
            new_a_states = []
            for a_state, action in zip(agent_states, actions):
                if a_state is None:
                    # 为什么智能体状态为空时可以构建一个维度和动作相同，且全为零的状态
                    # 这是针对初始状态，初始情况下，智能体的状态为空，则创建一个为0的状态
                    a_state = np.zeros(shape=action.shape, dtype=np.float32)
                # 书p298也 todo 了解OU过程
                a_state += self.ou_teta * (self.ou_mu - a_state)
                a_state += self.ou_sigma * np.random.normal(size=action.shape)

                action += self.ou_epsilon * a_state
                new_a_states.append(a_state)
        else:
            new_a_states = agent_states

        # 这个步骤 的意思是修正动作值到-1和1之间，否则将导致PyBullet抛出异常
        actions = np.clip(actions, -1, 1)
        return actions, new_a_states


class AgentD4PG(ptan.agent.BaseAgent):
    """
    Agent implementing noisy agent
    实现一个具备噪音的代理器
    """
    def __init__(self, net, device="cpu", epsilon=0.3):
        self.net = net
        self.device = device
        self.epsilon = epsilon

    def __call__(self, states, agent_states):
        '''
        states: 环境的状态
        agent_states: 智能体代理器的状态，在这里不需要关注这个状态

        return 执行的动作值，智能体的状态
        '''
        states_v = ptan.agent.float32_preprocessor(states).to(self.device)
        # 预测执行的动作值
        mu_v = self.net(states_v)
        actions = mu_v.data.cpu().numpy()
        # 给动作值添加噪音 todo 这里可以直接这么简单的+添加噪音吗？
        actions += self.epsilon * np.random.normal(size=actions.shape)
        # 防止动作值超过边界
        actions = np.clip(actions, -1, 1)
        return actions, agent_states


class LinearBottleNeck(nn.Module):

    def __init__(self, in_channels, out_channels, stride, t=6, class_num=100):
        super().__init__()

        self.residual = nn.Sequential(
            nn.Conv2d(in_channels, in_channels * t, 1),
            nn.BatchNorm2d(in_channels * t),
            nn.ReLU6(inplace=True),

            nn.Conv2d(in_channels * t, in_channels * t, 3, stride=stride, padding=1, groups=in_channels * t),
            nn.BatchNorm2d(in_channels * t),
            nn.ReLU6(inplace=True),

            nn.Conv2d(in_channels * t, out_channels, 1),
            nn.BatchNorm2d(out_channels)
        )

        self.stride = stride
        self.in_channels = in_channels
        self.out_channels = out_channels

    def forward(self, x):

        residual = self.residual(x)

        if self.stride == 1 and self.in_channels == self.out_channels:
            residual += x

        return residual

class ModelActor(nn.Module):

    def __init__(self, obs_size, act_size):
        super().__init__()
        obs = (obs_size[2], obs_size[0], obs_size[1])

        self.conv = nn.Sequential(
            nn.Conv2d(obs[0], 32, kernel_size=8, stride=4),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1),
            nn.ReLU()
        )
        conv_out_size = self._get_conv_out(obs)
        self.fc1 = nn.Sequential(
            nn.Linear(conv_out_size, 512),
            nn.ReLU(),
            nn.Linear(512, act_size))

        self.logstd = nn.Parameter(torch.zeros(act_size))


    def _get_conv_out(self, shape):
        o = self.conv(torch.zeros(1, *shape))
        return int(np.prod(o.size()))


    def forward(self, x):
        x = x.float() / 256
        x = self.conv(x)
        x = self.fc1(x.view(x.shape[0], -1))
        return x


class ModelCritic(nn.Module):

    def __init__(self, obs_size):
        super().__init__()

        obs = (obs_size[2], obs_size[0], obs_size[1])

        self.conv = nn.Sequential(
            nn.Conv2d(obs[0], 32, kernel_size=8, stride=4),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1),
            nn.ReLU()
        )
        conv_out_size = self._get_conv_out(obs)
        self.fc1 = nn.Sequential(
            nn.Linear(conv_out_size, 512),
            nn.ReLU(),
            nn.Linear(512, 1))

    def _get_conv_out(self, shape):
        o = self.conv(torch.zeros(1, *shape))
        return int(np.prod(o.size()))

    def forward(self, x):
        x = x.float() / 256
        x = self.conv(x)
        x = self.fc1(x.view(x.shape[0], -1))
        return x
    

class ModelActorLinear(nn.Module):
    def __init__(self, obs_size, act_size):
        '''
        :param obs_size: 观测的环境维度
        :param act_size: 执行的动作的维度
        '''
        super(ModelActorLinear, self).__init__()

        self.mu = nn.Sequential(
            nn.Linear(obs_size, HID_SIZE),
            nn.Tanh(),
            nn.Linear(HID_SIZE, HID_SIZE),
            nn.Tanh(),
            nn.Linear(HID_SIZE, act_size),
            nn.Tanh(),
        )
        # 作用 看名称像是方差：是用来控制动作的探索程度的
        # 怎么更新？：在训练的过程中，会不断的更新这个参数，更新的逻辑就在于计算熵损失以及计算动作优势大小的时候会参与计算，然后在梯度更新的时候，会自动更新这个参数到合适的大小
        self.logstd = nn.Parameter(torch.zeros(act_size))

    def forward(self, x):
        return self.mu(x)


class ModelCriticLinear(nn.Module):
    '''
    trop信赖域策略优化评价网络
    ACKTR算法中使用的critic网络
    ppt优化评价网络
    '''
    def __init__(self, obs_size):
        super(ModelCriticLinear, self).__init__()

        self.value = nn.Sequential(
            nn.Linear(obs_size, HID_SIZE),
            nn.ReLU(),
            nn.Linear(HID_SIZE, HID_SIZE),
            nn.ReLU(),
            nn.Linear(HID_SIZE, 1),
        )

    def forward(self, x):
        return self.value(x)


class MBv2BaseModel(nn.Module):

    def _make_stage(self, repeat, in_channels, out_channels, stride, t):

        layers = []
        layers.append(LinearBottleNeck(in_channels, out_channels, stride, t))

        while repeat - 1:
            layers.append(LinearBottleNeck(out_channels, out_channels, 1, t))
            repeat -= 1

        return nn.Sequential(*layers)


class DDPGActorMBv2(MBv2BaseModel):
    '''
    深度确定性策略梯度动作预测网络
    '''
    def __init__(self, obs_size, act_size):
        '''
        obs_size: 环境的维度
        act_size: 能够同时执行动作的个数（比如有多个手，不是每个手可以执行哪些动作）
        '''
        super(DDPGActorMBv2, self).__init__()

        obs_action = (obs_size[2], obs_size[0], obs_size[1])
        self.pre = nn.Sequential(
            nn.Conv2d(obs_action[0], 32, 1, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU6(inplace=True)
        )

        self.stage1 = LinearBottleNeck(32, 16, 1, 1)
        self.stage2 = self._make_stage(2, 16, 24, 2, 6)
        self.stage3 = self._make_stage(3, 24, 32, 2, 6)
        self.stage4 = self._make_stage(4, 32, 64, 2, 6)
        self.stage5 = self._make_stage(3, 64, 96, 1, 6)
        self.stage6 = self._make_stage(3, 96, 160, 1, 6)
        self.stage7 = LinearBottleNeck(160, 320, 1, 6)

        self.conv1 = nn.Sequential(
            nn.Conv2d(320, 1280, 1),
            nn.BatchNorm2d(1280),
            nn.ReLU6(inplace=True)
        )

        self.conv2 = nn.Conv2d(1280, act_size, 1)

    def forward(self, x):
        fx = x.float() / 256
        fx = self.pre(fx)
        fx = self.stage1(fx)
        fx = self.stage2(fx)
        fx = self.stage3(fx)
        fx = self.stage4(fx)
        fx = self.stage5(fx)
        fx = self.stage6(fx)
        fx = self.stage7(fx)
        fx = self.conv1(fx)
        fx = F.adaptive_avg_pool2d(fx, 1)
        fx = self.conv2(fx)
        fx = fx.view(fx.size(0), -1)

        return fx



class DDPGCriticMBv2(MBv2BaseModel):
    '''
    深度确定性策略梯度网络Q值评价网络
    '''
    def __init__(self, obs_size, act_size):
        '''
        obs_size: 环境的维度
        act_size: 能够同时执行动作的个数（比如有多个手，不是每个手可以执行哪些动作）
        '''
        super(DDPGCriticMBv2, self).__init__()

        obs_action = (obs_size[2], obs_size[0], obs_size[1])
        self.pre = nn.Sequential(
            nn.Conv2d(obs_action[0], 32, 1, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU6(inplace=True)
        )

        self.stage1 = LinearBottleNeck(32, 16, 1, 1)
        self.stage2 = self._make_stage(2, 16, 24, 2, 6)
        self.stage3 = self._make_stage(3, 24, 32, 2, 6)
        self.stage4 = self._make_stage(4, 32, 64, 2, 6)
        self.stage5 = self._make_stage(3, 64, 96, 1, 6)
        self.stage6 = self._make_stage(3, 96, 160, 1, 6)
        self.stage7 = LinearBottleNeck(160, 320, 1, 6)

        self.conv1 = nn.Sequential(
            nn.Conv2d(320, 1280, 1),
            nn.BatchNorm2d(1280),
            nn.ReLU6(inplace=True)
        )

        self.conv2 = nn.Conv2d(1280 + act_size, 1, 1)


    def forward(self, x, a):
        fx = x.float() / 256
        fx = self.pre(fx)
        fx = self.stage1(fx)
        fx = self.stage2(fx)
        fx = self.stage3(fx)
        fx = self.stage4(fx)
        fx = self.stage5(fx)
        fx = self.stage6(fx)
        fx = self.stage7(fx)
        fx = self.conv1(fx)
        fx = F.adaptive_avg_pool2d(fx, 1)
        a = a.unsqueeze(-1).unsqueeze(-1)
        fx = self.conv2(torch.cat((fx, a), dim=1))
        fx = fx.view(fx.size(0), -1)

        return fx



def bottle(f, x_tuple):
    x_sizes = tuple(map(lambda x: x.size(), x_tuple))
    y = f(*map(lambda x: x[0].contiguous().view(x[1][0] * x[1][1], *x[1][2:]), zip(x_tuple, x_sizes)))
    y_size = y.size()
    output = y.view(x_sizes[0][0], x_sizes[0][1], *y_size[1:])
    return output


class FeatureExtractor(nn.Module):
    def __init__(self, input_shape):
        super(FeatureExtractor, self).__init__()

        self.conv = nn.Sequential(
            nn.Conv2d(input_shape[0], 32, kernel_size=8, stride=4),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1),
            nn.ReLU()
        )


    def get_conv_out(self, shape):
        o = self.conv(torch.zeros(1, *shape))
        return int(np.prod(o.size()))

    def forward(self, x):
        x = x.float() / 255.0
        return self.conv(x)
    

class Dreamer(nn.Module):
    '''
    深度确定性策略梯度动作预测网络
    '''
    def __init__(self, belief_size, state_size, act_size, hidden_size, embedding_size, min_std_dev=0.1):
        '''
        obs_size: 环境的维度
        act_size: 能够同时执行动作的个数（比如有多个手，不是每个手可以执行哪些动作）
        '''
        super(Dreamer, self).__init__()

        self.min_std_dev = min_std_dev

        self.fc_embed_state_action = nn.Linear(state_size + act_size, belief_size)
        self.rnn = nn.GRUCell(belief_size, belief_size)
        self.fc_embed_belief_prior = nn.Linear(belief_size, hidden_size)
        self.fc_state_prior = nn.Linear(hidden_size, 2 * state_size)
        self.fc_embed_belief_posterior = nn.Linear(belief_size + embedding_size, hidden_size)
        self.fc_state_priorior = nn.Linear(hidden_size, 2 * state_size)



    def forward(self, prev_state, actions, prev_belief, observations: Optional[torch.Tensor] = None, nonterminals: Optional[torch.Tensor] = None):
        T = actions.size(0) + 1
        beliefs, prior_states, prior_means, prior_std_devs, posterior_states, posterior_means, posterior_std_devs = (
            [torch.empty(0)] * T,
            [torch.empty(0)] * T,
            [torch.empty(0)] * T,
            [torch.empty(0)] * T,
            [torch.empty(0)] * T,
            [torch.empty(0)] * T,
            [torch.empty(0)] * T
        )

        beliefs[0], prior_states[0], posterior_states[0] = prev_belief, prev_state, prev_state

        for t in range(T - 1):
            _state = (
                prior_states[t] if observations is None else posterior_states[t]
            )
            _state = (
                _state if nonterminals is None else _state * nonterminals[t]
            )

            hidden = F.elu(self.fc_embed_state_action(torch.cat([_state, actions[t]], dim=1)))
            beliefs[t+1] = self.rnn(hidden, beliefs[t])
            hidden = F.elu(self.fc_embed_belief_prior(beliefs[t+1]))
            prior_means[t + 1], _prior_std_dev = torch.chunk(self.fc_state_prior(hidden), 2, dim=1)
            prior_std_devs[t + 1] = F.softplus(_prior_std_dev) + self.min_std_dev
            prior_states[t + 1] = prior_means[t + 1] + prior_std_devs[t + 1] * torch.randn_like(prior_means[t + 1])
            if observations is not None:
                t_ = t - 1
                hidden = F.elu(self.fc_embed_belief_posterior(torch.cat([beliefs[t + 1], observations[t_ + 1]], dim=1)))
                posterior_means[t + 1], _posterior_std_dev = torch.chunk(self.fc_state_priorior(hidden), 2, dim=1)
                posterior_std_devs[t + 1] = F.softplus(_posterior_std_dev) + self.min_std_dev
                posterior_states[t + 1] = posterior_means[t + 1] + posterior_std_devs[t + 1] * torch.randn_like(posterior_means[t + 1])

        hidden = [
            torch.stack(beliefs[1:], dim=0),
            torch.stack(prior_states[1:], dim=0),
            torch.stack(prior_means[1:], dim=0),
            torch.stack(prior_std_devs[1:], dim=0),
        ]

        if observations is not None:
            hidden += [
                torch.stack(posterior_states[1:], dim=0),
                torch.stack(posterior_means[1:], dim=0),
                torch.stack(posterior_std_devs[1:], dim=0)
            ]

        return hidden
    
    

class ConvDreamerObservation(nn.Module):

    def __init__(self, belief_size, state_size, embedding_size):
        super().__init__()

        self.fc1 = nn.Linear(belief_size + state_size, embedding_size)
        self.embedding_size = embedding_size
        self.conv_decoder = nn.Sequential(
            nn.ConvTranspose2d(embedding_size, 128, kernel_size=7, stride=1),
            nn.ELU(),
            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),
            nn.ELU(),
            nn.ConvTranspose2d(64, 32, kernel_size=4, stride=2, padding=1),
            nn.ELU(),
            nn.ConvTranspose2d(32, 4, kernel_size=4, stride=3, padding=1, output_padding=1),
        )


    def forward(self, belief, state):
        x = torch.cat([belief, state], dim=1)
        x = self.fc1(x)
        x = x.view(-1, self.embedding_size, 1, 1)
        return self.conv_decoder(x)


class RewardDreamerModel(nn.Module):

    def __init__(self, belief_size, state_size, hidden_size):
        super().__init__()

        self.net = nn.Sequential(
            nn.Linear(belief_size + state_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, 1)
        )

    
    def forward(self, belief, state):
        return self.net(torch.cat([belief, state], dim=1)).squeeze(1)
    

class ConvDreamerEncoder(nn.Module):

    def __init__(self, obs_size, embedding_size):
        super().__init__()

        self.embedding_size = embedding_size
        self.conv1 = nn.Sequential(
            nn.Conv2d(obs_size[0], 32, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 128, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(128, 256, kernel_size=4, stride=2),
            nn.ReLU()
        )

        self.fc = nn.Identity() if embedding_size == 2304 else nn.Linear(2304, embedding_size)


    def forward(self, x):
        x = self.conv1(x)
        # 4 * 84 * 84 -> 256, 3, 3 = 2304
        x = x.view(-1, 2304)
        return self.fc(x)
    


class DreamerActorModel(nn.Module):

    def __init__(self, belief_size, state_size, hidden_size, act_size, dist='tanh_normal', min_std=1e-4, init_std=5, mean_scale=5, device='cpu'):
        super().__init__()

        self.net = nn.Sequential(
            nn.Linear(belief_size + state_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, 2 * act_size)
        )

        self._dist = dist
        self._min_std = torch.tensor(min_std, dtype=torch.float).to(device=device).detach()
        self._init_std = torch.tensor(init_std, dtype=torch.float).to(device=device).detach()
        self._mean_scale = torch.tensor(mean_scale, dtype=torch.float).to(device=device).detach()

    
    def forward(self, belief, state):
        raw_init_std = torch.log(torch.exp(self._init_std) - 1)
        x = torch.cat([belief, state], dim=1)
        action = self.net(x).squeeze(dim=1)

        action_mean, action_std_dev = torch.chunk(action, 2, dim=1)
        action_mean = self._mean_scale * torch.tanh(action_mean / self._mean_scale)
        action_std = F.softplus(action_std_dev + raw_init_std) + self._min_std

        return action_mean, action_std
    

    def get_action(self, belief, state, det=False):
        action_mean, action_std = self.forward(belief, state)
        dist = torch.distributions.normal.Normal(action_mean, action_std)
        dist = torch.distributions.transformed_distribution.TransformedDistribution(dist, TanhBijector())
        dist = torch.distributions.Independent(dist, 1)
        dist = SampleDist(dist)

        if det:
            return dist.mode()
        else:
            return dist.rsample()
        

class DreamerValueModel(nn.Module):
    def __init__(self, belief_size, state_size, hidden_size):
        super().__init__()

        self.net = nn.Sequential(
            nn.Linear(belief_size + state_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, 1)
        )


    def forward(self, belief, state):
        return self.net(torch.cat([belief, state], dim=1)).squeeze(1)


class SampleDist:

    def __init__(self, dist, samples=100):
        self._dist = dist
        self._samples = samples

    
    @property
    def name(self):
        return 'SampleDist'
    

    def __getattr__(self, name):
        return getattr(self._dist, name)
    

    def mean(self):
        sample = self._dist.rsample()
        return torch.mean(sample, dim=0)
    

    def mode(self):
        dist = self._dist.expand((self._samples, *self._dist.batch_shape))
        sample = dist.rsample()
        logprob = dist.log_prob(sample)
        batch_size = sample.size(1)
        feature_size = sample.size(2)
        indices = torch.argmax(logprob, dim=0).reshape(1, batch_size, 1).expand(1, batch_size, feature_size)
        return torch.gather(sample, 0, indices).squeeze(0)
    

    def entropy(self):
        dist = self._dist.expand((self._samples, *self._dist.batch_shape))
        sample = dist.rsample()
        logprob = dist.log_prob(sample)
        return -torch.mean(logprob, dim=0)
    

    def sample(self):
        return self._dist.sample()
        

class TanhBijector(torch.distributions.Transform):

    def __init__(self):
        super().__init__()

        self.bijective = True
        self.domain = torch.distributions.constraints.real
        self.codomain = torch.distributions.constraints.interval(-1.0, 1.0)

    
    @property
    def sign(self):
        return 1.0
    

    @staticmethod
    def atanh(x):
        return 0.5 * torch.log((1 + x) / (1 - x))
    

    def _call(self, x):
        return torch.tanh(x)
    

    def _inverse(self, y: torch.Tensor):
        y = torch.where((torch.abs(y) <= 1.0), torch.clamp(y, -0.99999997, 0.99999997), y)
        y = TanhBijector.atanh(y)
        return y
    

    def log_abs_det_jacobian(self, x, y):
        return 2.0 * (np.log(2) - x - F.softplus(-2.0 * x))


class PixelEncoder(nn.Module):
    # for 84 x 84 inputs
    OUT_DIM = {2: 39, 4: 35, 6: 31}
    # for 64 x 64 inputs
    OUT_DIM_64 = {2: 29, 4: 25, 6: 21}
    '''像素空间编码器'''
    """Convolutional encoder of pixels observations."""
    def __init__(self, obs_shape, feature_dim, num_layers=2, num_filters=32,output_logits=False):
        '''
        param output_logits: 在实际使用时均为True
        '''
        super().__init__()

        assert len(obs_shape) == 3
        self.obs_shape = obs_shape
        self.feature_dim = feature_dim
        self.num_layers = num_layers

        # 卷积层，没有激活函数？没有归一化层
        # 在forward中加入了relu
        self.convs = nn.ModuleList(
            [nn.Conv2d(obs_shape[0], num_filters, 3, stride=2)]
        )
        for i in range(num_layers - 1):
            self.convs.append(nn.Conv2d(num_filters, num_filters, 3, stride=1))

        # 接着就是线性层 + 层归一化
        out_dim = PixelEncoder.OUT_DIM_64[num_layers] if obs_shape[-1] == 64 else PixelEncoder.OUT_DIM[num_layers] 
        self.fc = nn.Linear(num_filters * out_dim * out_dim, self.feature_dim)
        self.ln = nn.LayerNorm(self.feature_dim)

        # outputs存储中间结果，比如obs，conv1，conv2，fc，ln，tanh
        self.outputs = dict()
        self.output_logits = output_logits

    def reparameterize(self, mu, logstd):
        std = torch.exp(logstd)
        eps = torch.randn_like(std)
        return mu + eps * std

    def forward_conv(self, obs):
        # 对环境进行归一化预处理
        obs = obs / 255.
        self.outputs['obs'] = obs

        # 进行卷积操作
        conv = torch.relu(self.convs[0](obs))
        self.outputs['conv1'] = conv

        for i in range(1, self.num_layers):
            conv = torch.relu(self.convs[i](conv))
            self.outputs['conv%s' % (i + 1)] = conv

        # 返回编码特征向量
        h = conv.view(conv.size(0), -1)
        return h

    def forward(self, obs, detach=False):
        '''
        params obs: 环境编码器
        params detach: 是否需要detach todo 作用 ，在eval时为false，没必要断开，因为不计算梯度
        params output_logits: 是否需要输出logits，决定输出结果是否经过tanh函数
        '''
        h = self.forward_conv(obs)

        if detach:
            # detach原因：在CURL中的应用场景Actor网络更新时不需要更新编码器
            # 目标编码器计算时需要固定参数
            # 防止编码器被过度更新
            # 稳定训练过程
            # 允许不同组件独立学习
            # todo 那为啥组合网络，一边仅优化actor 线性层 一边优化所有
            # todo 实现actor和critic合并
            h = h.detach() # 断开梯度传播

        h_fc = self.fc(h)
        self.outputs['fc'] = h_fc

        h_norm = self.ln(h_fc)
        self.outputs['ln'] = h_norm

        if self.output_logits:
            '''
            todo 实际使用时都是true，可能一般不使用
            直接输出层归一化(LayerNorm)后的值
            值域不受限制
            适用于需要原始特征分布的情况
            在CURL中用于对比学习
            '''
            out = h_norm
        else:
            '''
            todo
            对层归一化后的值进行tanh激活
            将值压缩到[-1, 1]区间
            适用于需要有界输出的情况
            可以防止特征值过大
            '''
            out = torch.tanh(h_norm)
            self.outputs['tanh'] = out

        return out

    def copy_conv_weights_from(self, source):
        """Tie convolutional layers"""
        # only tie conv layers
        for i in range(self.num_layers):
            common.tie_weights(src=source.convs[i], trg=self.convs[i])

    def log(self, L, step, log_freq):
        if step % log_freq != 0:
            return

        for k, v in self.outputs.items():
            L.log_histogram('train_encoder/%s_hist' % k, v, step)
            if len(v.shape) > 2:
                L.log_image('train_encoder/%s_img' % k, v[0], step)

        for i in range(self.num_layers):
            L.log_param('train_encoder/conv%s' % (i + 1), self.convs[i], step)
        L.log_param('train_encoder/fc', self.fc, step)
        L.log_param('train_encoder/ln', self.ln, step)


LOG_FREQ = 10000
class Actor(nn.Module):
    """MLP actor network."""
    def __init__(
        self, obs_shape, action_shape, hidden_dim, encoder_type,
        encoder_feature_dim, log_std_min, log_std_max, num_layers, num_filters
    ):
        super().__init__()

        # 创建编码器
        self.encoder = PixelEncoder(
            obs_shape, encoder_feature_dim, num_layers,
            num_filters, output_logits=True
        )

        self.log_std_min = log_std_min
        self.log_std_max = log_std_max

        # 创建header了，线性空间，有激活函数
        self.trunk = nn.Sequential(
            nn.Linear(self.encoder.feature_dim, hidden_dim), nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim), nn.ReLU(),
            nn.Linear(hidden_dim, 2 * action_shape[0])
        )

        # 存储中间值，有均值、方差、
        self.outputs = dict()
        self.apply(common.weight_init)


    @staticmethod
    def gaussian_logprob(noise, log_std):
        """Compute Gaussian log probability."""
        residual = (-0.5 * noise.pow(2) - log_std).sum(-1, keepdim=True)
        return residual - 0.5 * np.log(2 * np.pi) * noise.size(-1)
    

    @staticmethod
    def squash(mu, pi, log_pi):
        """Apply squashing function.
        See appendix C from https://arxiv.org/pdf/1812.05905.pdf. todo
        函数实现了 SAC 算法中的 tanh 压缩操作，其作用是将无界的动作值压缩到有界区间。让我分析一下这个函数
        """
        # 压缩到[-1, 1]区间
        mu = torch.tanh(mu)
        # 对采样的动作进行压缩
        if pi is not None:
            pi = torch.tanh(pi)
        # 对数概率进行修正
        if log_pi is not None:
            # 根据变量替换法则修正对数概率
            log_pi -= torch.log(F.relu(1 - pi.pow(2)) + 1e-6).sum(-1, keepdim=True)
        return mu, pi, log_pi


    def forward(
        self, obs, compute_pi=True, compute_log_pi=True, detach_encoder=False
    ):
        '''
        obs: 观察值

        环境的观察输入(状态)
        通常是图像数据
        compute_pi=True: 是否计算策略动作,

        True: 计算随机策略动作(用于训练和探索)，在eval时为True
        False: 只返回确定性动作(用于评估)
        compute_log_pi=True: 是否计算对数概率

        True: 计算动作的对数概率(用于计算策略损失)
        False: 不计算对数概率,在eval时为false
        detach_encoder=False: 是否分离编码器梯度

        True: 停止梯度传播到编码器
        False: 允许梯度传播到编码器 在eval时为false
        '''
        # 环境编码器得到的时均值和方差
        obs = self.encoder(obs, detach=detach_encoder)

        # 拆分预测值分别为均值和方差
        mu, log_std = self.trunk(obs).chunk(2, dim=-1)

        # constrain log_std inside [log_std_min, log_std_max]
        # 限制方差的大小，可能是为防止抖动把 todo
        '''
        是的，这段代码是为了限制方差的大小。让我解析一下这段代码：

        ```python
        # constrain log_std inside [log_std_min, log_std_max]
        log_std = torch.tanh(log_std)  # 首先压缩到(-1,1)范围
        log_std = self.log_std_min + 0.5 * (self.log_std_max - self.log_std_min) * (log_std + 1)
        ```

        这个限制操作分两步：

        1. **首先使用tanh函数压缩**:
        - `torch.tanh(log_std)` 将输入压缩到(-1,1)范围
        - tanh函数可以防止数值过大或过小

        2. **然后线性映射到指定范围**:
        - `log_std_min` 默认为-10
        - `log_std_max` 默认为2
        - `(log_std + 1)` 将范围调整到(0,2)
        - `0.5 * (self.log_std_max - self.log_std_min)` 计算目标范围的一半
        - 最终将值映射到[log_std_min, log_std_max]范围内

        限制方差的目的：
        1. 防止动作探索过度发散
        2. 保持策略的稳定性
        3. 控制随机性的程度
        4. 避免数值不稳定

        这样做可以让策略网络输出的动作分布保持在一个合理的范围内，既不会太确定(方差太小)也不会太随机(方差太大)。
        '''
        log_std = torch.tanh(log_std)
        log_std = self.log_std_min + 0.5 * (
            self.log_std_max - self.log_std_min
        ) * (log_std + 1)

        self.outputs['mu'] = mu
        self.outputs['std'] = log_std.exp()

        if compute_pi:
            # 对动作进行采样，使用的是高斯分布
            std = log_std.exp()
            noise = torch.randn_like(mu)
            pi = mu + noise * std
        else:
            pi = None
            entropy = None

        # 在eval时不计算对数概率
        if compute_log_pi:
            log_pi = Actor.gaussian_logprob(noise, log_std)
        else:
            log_pi = None

        mu, pi, log_pi = Actor.squash(mu, pi, log_pi)

        # 返回预测的均值、采样的动作、对数概率和标准差
        return mu, pi, log_pi, log_std

    def log(self, L, step, log_freq=LOG_FREQ):
        if step % log_freq != 0:
            return

        for k, v in self.outputs.items():
            L.log_histogram('train_actor/%s_hist' % k, v, step)

        L.log_param('train_actor/fc1', self.trunk[0], step)
        L.log_param('train_actor/fc2', self.trunk[2], step)
        L.log_param('train_actor/fc3', self.trunk[4], step)


class QFunction(nn.Module):
    """MLP for q-function."""
    def __init__(self, obs_dim, action_dim, hidden_dim):
        super().__init__()

        self.trunk = nn.Sequential(
            nn.Linear(obs_dim + action_dim, hidden_dim), nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim), nn.ReLU(),
            nn.Linear(hidden_dim, 1)
        )

    def forward(self, obs, action):
        assert obs.size(0) == action.size(0)

        obs_action = torch.cat([obs, action], dim=1)
        return self.trunk(obs_action)


class Critic(nn.Module):
    """Critic network, employes two q-functions."""
    def __init__(
        self, obs_shape, action_shape, hidden_dim, encoder_type,
        encoder_feature_dim, num_layers, num_filters
    ):
        super().__init__()

        # 环境编码器
        self.encoder = PixelEncoder(
            obs_shape, encoder_feature_dim, num_layers,
            num_filters, output_logits=True
        )

        # 两个线性编码器，作用是啥？对比我现在已有的SAC看看
        self.Q1 = QFunction(
            self.encoder.feature_dim, action_shape[0], hidden_dim
        )
        self.Q2 = QFunction(
            self.encoder.feature_dim, action_shape[0], hidden_dim
        )

        self.outputs = dict()
        self.apply(common.weight_init)

    def forward(self, obs, action, detach_encoder=False):
        # detach_encoder allows to stop gradient propogation to encoder
        obs = self.encoder(obs, detach=detach_encoder)

        q1 = self.Q1(obs, action)
        q2 = self.Q2(obs, action)

        self.outputs['q1'] = q1
        self.outputs['q2'] = q2

        return q1, q2

    def log(self, L, step, log_freq=LOG_FREQ):
        if step % log_freq != 0:
            return

        self.encoder.log(L, step, log_freq)

        for k, v in self.outputs.items():
            L.log_histogram('train_critic/%s_hist' % k, v, step)

        for i in range(3):
            L.log_param('train_critic/q1_fc%d' % i, self.Q1.trunk[i * 2], step)
            L.log_param('train_critic/q2_fc%d' % i, self.Q2.trunk[i * 2], step)


class CURL(nn.Module):
    """
    CURL 
    """

    def __init__(self, obs_shape, z_dim, batch_size, critic, critic_target, output_type="continuous"):
        '''
        obs_shape: 观察空间的形状

        用于处理环境状态的输入维度
        通常是图像数据,形状为 (channels, height, width)
        z_dim: 隐空间维度

        编码器输出的特征向量维度
        决定了表征学习的维度大小
        在这个实现中等于 encoder_feature_dim (默认50)
        batch_size: 批量大小

        训练时每批数据的样本数
        用于计算对比损失时的批次维度
        critic: Critic 网络实例

        包含主要的编码器网络
        CURL 会使用 critic.encoder 作为其编码器
        critic_target: 目标 Critic 网络实例

        包含目标编码器网络
        CURL 使用 critic_target.encoder 作为目标编码器
        用于计算对比学习的正样本编码
        output_type: 输出类型

        默认为 "continuous"
        指定输出空间类型
        这个实现中仅支持连续动作空间
        '''

        super(CURL, self).__init__()
        self.batch_size = batch_size

        # 评价网络的环境编码器
        self.encoder = critic.encoder

        self.encoder_target = critic_target.target_model.encoder 

        # todo W的参数
        self.W = nn.Parameter(torch.rand(z_dim, z_dim))
        # 输出的动作空间类型，todo 是否支持离散动作空间
        self.output_type = output_type

    def encode(self, x, detach=False, ema=False):
        """
        Encoder: z_t = e(x_t)
        :param x: x_t, x y coordinates
        :return: z_t, value in r2
        """
        if ema:
            with torch.no_grad():
                z_out = self.encoder_target(x)
        else:
            z_out = self.encoder(x)

        if detach:
            z_out = z_out.detach()
        return z_out

    def compute_logits(self, z_a, z_pos):
        """
        todo 这里计算了CURL损失，了解原理
        z_a：当前状态的编码。
        z_pos：目标状态的编码

        Uses logits trick for CURL:
        - compute (B,B) matrix z_a (W z_pos.T)
        - positives are all diagonal elements
        - negatives are all other elements
        - to compute loss use multiclass cross entropy with identity matrix for labels
        """
        Wz = torch.matmul(self.W, z_pos.T)  # (z_dim,B)
        logits = torch.matmul(z_a, Wz)  # (B,B)
        logits = logits - torch.max(logits, 1)[0][:, None]
        return logits
